import numpy as np
import csv

# 定义全局协议列表
protocol_list = ['tcp', 'udp', 'icmp']
service_list = ['aol', 'auth', 'bgp', 'courier', 'csnet_ns', 'ctf', 'daytime', 'discard', 'domain', 'domain_u',
                'echo', 'eco_i', 'ecr_i', 'efs', 'exec', 'finger', 'ftp', 'ftp_data', 'gopher', 'harvest', 'hostnames',
                'http', 'http_2784', 'http_443', 'http_8001', 'imap4', 'IRC', 'iso_tsap', 'klogin', 'kshell', 'ldap',
                'link', 'login', 'mtp', 'name', 'netbios_dgm', 'netbios_ns', 'netbios_ssn', 'netstat', 'nnsp', 'nntp',
                'ntp_u', 'other', 'pm_dump', 'pop_2', 'pop_3', 'printer', 'private', 'red_i', 'remote_job', 'rje',
                'shell',
                'smtp', 'sql_net', 'ssh', 'sunrpc', 'supdup', 'systat', 'telnet', 'tftp_u', 'tim_i', 'time', 'urh_i',
                'urp_i',
                'uucp', 'uucp_path', 'vmnet', 'whois', 'X11', 'Z39_50']
flag_list = ['OTH', 'REJ', 'RSTO', 'RSTOS0', 'RSTR', 'S0', 'S1', 'S2', 'S3', 'SF', 'SH']
label_list = ['normal.', 'buffer_overflow.', 'loadmodule.', 'perl.', 'neptune.', 'smurf.',
              'guess_passwd.', 'pod.', 'teardrop.', 'portsweep.', 'ipsweep.', 'land.', 'ftp_write.',
              'back.', 'imap.', 'satan.', 'phf.', 'nmap.', 'multihop.', 'warezmaster.', 'warezclient.',
              'spy.', 'rootkit.']  # 初始已知标签


# 取字符对应的索引
def find_index(x, y):
    return [i for i in range(len(y)) if y[i] == x]


def handleProtocol(inputs):
    if inputs[1] in protocol_list:
        return find_index(inputs[1], protocol_list)[0]
    return len(protocol_list)  # 处理未知协议


def handleService(inputs):
    if inputs[2] in service_list:
        return find_index(inputs[2], service_list)[0]
    return len(service_list)  # 处理未知服务


def handleFlag(inputs):
    if inputs[3] in flag_list:
        return find_index(inputs[3], flag_list)[0]
    return len(flag_list)  # 处理未知状态


def handleLabel(inputs):
    global label_list
    if inputs[41] in label_list:
        return find_index(inputs[41], label_list)[0]
    else:
        # 处理未知攻击类型
        label_list.append(inputs[41])
        return len(label_list) - 1  # 返回新添加的索引


if __name__ == '__main__':
    data_numerization = open("kddcup.data.numerization.txt", 'w', newline='')

    with open('kddcup.data.txt', 'r') as data_original:
        csv_reader = csv.reader(data_original)
        csv_writer = csv.writer(data_numerization, dialect='excel')

        for row in csv_reader:
            # 检查行长度是否足够
            if len(row) < 42:
                print(f"跳过不完整的行: {row}")
                continue

            temp_line = np.array(row)

            # 安全处理每个字段
            try:
                temp_line[1] = handleProtocol(row)
                temp_line[2] = handleService(row)
                temp_line[3] = handleFlag(row)
                temp_line[41] = handleLabel(row)
            except IndexError as e:
                print(f"处理行时出错: {row}\n错误: {e}")
                continue

            csv_writer.writerow(temp_line)

    data_numerization.close()
    print('数值化完成！')
    print(f"最终标签列表包含 {len(label_list)} 种攻击类型")